"""
GomerX 视觉处理模块
用于图像处理、线条检测和颜色识别
"""

# 导入必要的系统模块
import cv2  # OpenCV库，用于图像处理和计算机视觉
import numpy as np  # NumPy库，用于数值计算和数组操作
from config.settings import VISION_CONFIG  # 导入视觉处理配置参数

class VisionProcessor:
    """GomerX 视觉处理类
    
    负责图像预处理、线条检测、颜色识别和物体检测等视觉处理任务
    提供绘制检测结果和计算位置偏移的功能
    """
    
    def __init__(self, config=None):
        """初始化视觉处理器
        
        从配置文件加载视觉处理参数，包括线条检测、颜色识别和物体检测的配置
        
        Args:
            config: 视觉处理配置字典，如果为None则使用默认配置
        """
        # 如果提供了配置，则使用提供的配置，否则使用默认配置
        if config is None:
            # 从配置文件加载各类视觉处理参数
            self.line_config = VISION_CONFIG["line_detection"]  # 线条检测配置参数
            self.color_config = VISION_CONFIG["color_detection"]  # 颜色识别配置参数
            self.object_config = VISION_CONFIG["object_detection"]  # 物体检测配置参数
            self.line_following_config = VISION_CONFIG.get("line_following", {})  # 巡线配置参数
        else:
            # 从提供的配置中加载参数
            self.line_config = config.get("line_detection", {})
            self.color_config = config.get("color_detection", {})
            self.object_config = config.get("object_detection", {})
            self.line_following_config = config.get("line_following", {})
    
    def preprocess_image(self, image):
        """
        图像预处理
        
        对输入图像进行灰度转换和高斯模糊处理，为后续的边缘检测做准备
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            
        Returns:
            numpy.ndarray: 预处理后的灰度图像
        """
        # 将BGR图像转换为灰度图像
        # 灰度图像只有一个通道，可以简化计算并提高处理速度
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊对图像进行去噪处理
        # (5, 5)是高斯核的大小，0表示自动计算标准差
        # 高斯模糊可以平滑图像，减少噪声对边缘检测的影响
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        return blurred
    
    def detect_lines(self, image):
        """
        检测图像中的线条
        
        使用Canny边缘检测和霍夫变换检测图像中的直线
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            
        Returns:
            list: 检测到的线条列表，每条线表示为 [x1, y1, x2, y2]
        """
        # 对图像进行预处理，转换为灰度并去噪
        processed = self.preprocess_image(image)
        
        # 使用Canny算法进行边缘检测
        # Canny边缘检测是一种多级边缘检测算法，能够检测图像中的强边缘
        # canny_low和canny_high是Canny算法的两个阈值，用于边缘连接
        edges = cv2.Canny(
            processed, 
            self.line_config["canny_low"],  # 低阈值
            self.line_config["canny_high"]  # 高阈值
        )
        
        # 使用概率霍夫变换检测直线
        # 霍夫变换是一种用于检测图像中几何形状的技术
        # HoughLinesP是概率霍夫变换，相比标准霍夫变换更高效
        lines = cv2.HoughLinesP(
            edges,  # 输入边缘图像
            rho=1,  # 距离分辨率（像素）
            theta=np.pi / 180,  # 角度分辨率（弧度）
            threshold=self.line_config["threshold"],  # 累加器阈值，越大检测的线越少
            minLineLength=self.line_config["min_line_length"],  # 最小线长度
            maxLineGap=self.line_config["max_line_gap"]  # 最大线间隙
        )
        
        # 如果没有检测到线条，返回空列表
        if lines is None:
            return []
            
        # 将线条格式转换为列表形式
        result_lines = []
        for line in lines:
            x1, y1, x2, y2 = line[0]  # 提取线条的起点和终点坐标
            result_lines.append([x1, y1, x2, y2])  # 添加到结果列表
            
        return result_lines
    
    def calculate_line_offset(self, lines, image_width):
        """
        计算线条相对于图像中心的偏移量
        
        计算所有线条的平均位置相对于图像中心的偏移量，用于循迹控制
        
        Args:
            lines: 检测到的线条列表，每条线表示为 [x1, y1, x2, y2]
            image_width: 图像宽度（像素）
            
        Returns:
            float: 归一化偏移量，范围[-1, 1]
                   正值表示线条在右侧，负值表示在左侧，0表示居中
        """
        # 如果没有检测到线条，返回0表示居中
        if not lines:
            return 0
            
        # 计算所有线条的中点
        midpoints = []
        for line in lines:
            x1, y1, x2, y2 = line
            # 计算线条中点的x坐标
            midpoint_x = (x1 + x2) / 2
            midpoints.append(midpoint_x)
            
        # 计算所有线条中点的平均值
        avg_midpoint = sum(midpoints) / len(midpoints)
        
        # 计算相对于图像中心的偏移量
        image_center = image_width / 2  # 图像中心的x坐标
        offset = avg_midpoint - image_center  # 实际偏移量（像素）
        
        # 将偏移量归一化到[-1, 1]范围
        # 这样做可以使偏移量与图像分辨率无关，提高算法的通用性
        normalized_offset = offset / (image_width / 2)
        
        return normalized_offset
    
    def detect_color(self, image, color_name):
        """
        检测图像中特定颜色的区域
        
        使用HSV色彩空间和颜色阈值检测图像中的特定颜色区域
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            color_name: 颜色名称，如 "red", "blue", "green"
            
        Returns:
            tuple: (mask, contours) 颜色掩码和轮廓列表
                   mask是二值图像，白色区域表示检测到的颜色
                   contours是检测到的颜色区域的轮廓列表
        """
        # 检查颜色名称是否在配置中
        if color_name not in self.color_config:
            raise ValueError(f"不支持的颜色: {color_name}")
            
        # 获取颜色范围的HSV阈值
        color_range = self.color_config[color_name]
        lower = np.array(color_range["lower"])  # HSV下限值
        upper = np.array(color_range["upper"])  # HSV上限值
        
        # 将BGR图像转换为HSV色彩空间
        # HSV色彩空间更适合颜色检测，因为它将颜色信息（H）与亮度信息（V）分离
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 创建颜色掩码
        # inRange函数创建一个二值图像，其中在指定HSV范围内的像素为白色（255），其他为黑色（0）
        mask = cv2.inRange(hsv, lower, upper)
        
        # 使用形态学操作去除噪点
        # 创建5x5的矩形结构元素
        kernel = np.ones((5, 5), np.uint8)
        
        # 开运算：先腐蚀后膨胀，去除小的白色噪点
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        # 闭运算：先膨胀后腐蚀，填充小的黑色空洞
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 查找轮廓
        # RETR_EXTERNAL：只检测最外层轮廓
        # CHAIN_APPROX_SIMPLE：压缩轮廓，节省内存
        contours, _ = cv2.findContours(
            mask,  # 输入二值图像
            cv2.RETR_EXTERNAL,  # 轮廓检索模式
            cv2.CHAIN_APPROX_SIMPLE  # 轮廓近似方法
        )
        
        return mask, contours
    
    def filter_objects(self, contours):
        """
        根据面积过滤轮廓，保留可能是物体的轮廓
        
        过滤掉过小或过大的轮廓，保留可能是实际物体的轮廓
        
        Args:
            contours: 轮廓列表，每个轮廓是一个点集
            
        Returns:
            list: 过滤后的轮廓列表
        """
        filtered = []  # 存储过滤后的轮廓
        
        # 从配置中获取面积阈值
        min_area = self.object_config["min_contour_area"]  # 最小面积阈值
        max_area = self.object_config["max_contour_area"]  # 最大面积阈值
        
        # 遍历所有轮廓
        for contour in contours:
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            
            # 检查面积是否在有效范围内
            if min_area <= area <= max_area:
                filtered.append(contour)  # 添加到过滤后的列表
                
        return filtered
    
    def get_object_position(self, contour, image_width):
        """
        获取物体相对于图像中心的位置
        
        计算物体的边界框和相对于图像中心的偏移量
        
        Args:
            contour: 物体轮廓
            image_width: 图像宽度（像素）
            
        Returns:
            tuple: (x, y, width, height, offset) 
                   x, y: 物体边界框左上角坐标
                   width, height: 物体边界框的宽度和高度
                   offset: 归一化偏移量，范围[-1, 1]
        """
        # 获取轮廓的边界矩形
        x, y, w, h = cv2.boundingRect(contour)
        
        # 计算物体中心的x坐标
        center_x = x + w / 2
        
        # 计算相对于图像中心的偏移量
        image_center = image_width / 2  # 图像中心的x坐标
        offset = center_x - image_center  # 实际偏移量（像素）
        
        # 将偏移量归一化到[-1, 1]范围
        normalized_offset = offset / (image_width / 2)
        
        return x, y, w, h, normalized_offset
    
    def draw_detection_results(self, image, lines=None, objects=None):
        """
        在图像上绘制检测结果
        
        将检测到的线条和物体绘制在图像上，用于可视化调试
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            lines: 检测到的线条列表，每条线表示为 [x1, y1, x2, y2]
            objects: 检测到的物体列表，每个元素为 (contour, color_name)
            
        Returns:
            numpy.ndarray: 绘制了检测结果的图像
        """
        # 创建图像副本，避免修改原始图像
        result = image.copy()
        
        # 绘制线条
        if lines:
            for line in lines:
                x1, y1, x2, y2 = line  # 提取线条的起点和终点坐标
                # 在图像上绘制绿色线条，线宽为2像素
                cv2.line(result, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # 绘制物体
        if objects:
            for contour, color_name in objects:
                # 绘制轮廓
                # 使用红色绘制轮廓，线宽为2像素
                cv2.drawContours(result, [contour], -1, (0, 0, 255), 2)
                
                # 获取边界框
                x, y, w, h = cv2.boundingRect(contour)
                
                # 绘制边界框
                # 使用蓝色绘制边界框，线宽为2像素
                cv2.rectangle(result, (x, y), (x+w, y+h), (255, 0, 0), 2)
                
                # 添加颜色标签
                # 在边界框上方添加颜色名称文本
                cv2.putText(result, color_name, (x, y-10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        return result
    
    # ==================== 巡线检测功能增强 ====================
    
    def detect_line_for_following(self, image):
        """
        专门用于巡线的线条检测
        
        使用增强的图像处理技术，针对巡线场景优化线条检测算法
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            
        Returns:
            dict: 包含线条检测结果的字典
                - 'lines': 检测到的线条列表
                - 'binary': 二值化图像
                - 'roi': 感兴趣区域图像
                - 'offset': 相对于中心的偏移量
                - 'angle': 线条角度
                - 'status': 巡线状态 ('normal', 'lost', 'intersection', 'sharp_turn')
        """
        # 获取图像尺寸
        height, width = image.shape[:2]
        
        # 定义感兴趣区域(ROI)，通常是图像的下半部分
        roi_ratio = self.line_following_config.get("roi_ratio", 0.6)
        roi_top = int(height * (1 - roi_ratio))
        roi = image[roi_top:height, :]
        
        # 转换为灰度图像
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        
        # 自适应阈值二值化，比固定阈值更能适应光照变化
        block_size = self.line_following_config.get("adaptive_block_size", 15)
        c_value = self.line_following_config.get("adaptive_c", 5)
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, block_size, c_value
        )
        
        # 形态学操作，去除噪点并连接断开的线段
        kernel_size = self.line_following_config.get("morph_kernel_size", 3)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        # 使用霍夫变换检测线条
        lines = cv2.HoughLinesP(
            binary, 1, np.pi/180, 
            threshold=self.line_following_config.get("hough_threshold", 20),
            minLineLength=self.line_following_config.get("min_line_length", 20),
            maxLineGap=self.line_following_config.get("max_line_gap", 10)
        )
        
        # 处理检测结果
        result = {
            'lines': [],
            'binary': binary,
            'roi': roi,
            'offset': 0,
            'angle': 0,
            'status': 'lost'
        }
        
        if lines is not None and len(lines) > 0:
            # 调整线条坐标到原始图像坐标系
            adjusted_lines = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                # 将y坐标调整回原始图像坐标系
                y1 += roi_top
                y2 += roi_top
                adjusted_lines.append([x1, y1, x2, y2])
            
            result['lines'] = adjusted_lines
            
            # 计算偏移量
            result['offset'] = self.calculate_line_offset(adjusted_lines, width)
            
            # 计算平均角度
            angles = []
            for x1, y1, x2, y2 in adjusted_lines:
                angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi
                angles.append(angle)
            
            if angles:
                result['angle'] = np.mean(angles)
            
            # 判断巡线状态
            result['status'] = self._determine_line_status(adjusted_lines, width, height)
        
        return result
    
    def _determine_line_status(self, lines, width, height):
        """
        根据检测到的线条判断巡线状态
        
        Args:
            lines: 检测到的线条列表
            width: 图像宽度
            height: 图像高度
            
        Returns:
            str: 巡线状态 ('normal', 'lost', 'intersection', 'sharp_turn')
        """
        if not lines:
            return 'lost'
        
        # 计算线条的平均角度
        angles = []
        for x1, y1, x2, y2 in lines:
            angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi
            angles.append(angle)
        
        avg_angle = np.mean(angles)
        
        # 判断是否为急转弯
        if abs(avg_angle) > self.line_following_config.get("sharp_turn_threshold", 30):
            return 'sharp_turn'
        
        # 判断是否为交叉路口
        # 如果检测到多条角度差异较大的线，可能是交叉路口
        if len(lines) >= 3:
            angle_diff = max(angles) - min(angles)
            if angle_diff > self.line_following_config.get("intersection_threshold", 45):
                return 'intersection'
        
        # 正常巡线状态
        return 'normal'
    
    def calculate_steering_angle(self, offset, line_status='normal'):
        """
        根据线条偏移量和状态计算转向角度
        
        Args:
            offset: 线条相对于中心的偏移量，范围[-1, 1]
            line_status: 巡线状态 ('normal', 'lost', 'intersection', 'sharp_turn')
            
        Returns:
            float: 转向角度（度），正值表示右转，负值表示左转
        """
        # 从配置中获取参数
        base_gain = self.line_following_config.get("steering_gain", 30.0)
        intersection_gain = self.line_following_config.get("intersection_gain", 15.0)
        sharp_turn_gain = self.line_following_config.get("sharp_turn_gain", 45.0)
        
        # 根据状态选择不同的增益
        if line_status == 'lost':
            # 线条丢失时，保持当前方向或小幅摆动寻找线条
            return 0
        elif line_status == 'intersection':
            # 交叉路口，减小转向幅度
            return offset * intersection_gain
        elif line_status == 'sharp_turn':
            # 急转弯，增加转向幅度
            return offset * sharp_turn_gain
        else:
            # 正常巡线
            return offset * base_gain
    
    def detect_lane_markings(self, image):
        """
        检测车道标记线（左线和右线）
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            
        Returns:
            dict: 包含左右车道线的检测结果
                - 'left_line': 左车道线参数 (x1, y1, x2, y2) 或 None
                - 'right_line': 右车道线参数 (x1, y1, x2, y2) 或 None
                - 'lane_center': 车道中心偏移量
                - 'lane_width': 车道宽度（像素）
        """
        # 获取图像尺寸
        height, width = image.shape[:2]
        
        # 定义感兴趣区域(ROI)
        roi_ratio = self.line_following_config.get("roi_ratio", 0.6)
        roi_top = int(height * (1 - roi_ratio))
        roi = image[roi_top:height, :]
        
        # 转换为灰度图像
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        
        # 自适应阈值二值化
        block_size = self.line_following_config.get("adaptive_block_size", 15)
        c_value = self.line_following_config.get("adaptive_c", 5)
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, block_size, c_value
        )
        
        # 使用霍夫变换检测线条
        lines = cv2.HoughLinesP(
            binary, 1, np.pi/180, 
            threshold=self.line_following_config.get("hough_threshold", 20),
            minLineLength=self.line_following_config.get("min_line_length", 20),
            maxLineGap=self.line_following_config.get("max_line_gap", 10)
        )
        
        result = {
            'left_line': None,
            'right_line': None,
            'lane_center': 0,
            'lane_width': 0
        }
        
        if lines is not None and len(lines) > 0:
            # 调整线条坐标到原始图像坐标系
            adjusted_lines = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                # 将y坐标调整回原始图像坐标系
                y1 += roi_top
                y2 += roi_top
                adjusted_lines.append([x1, y1, x2, y2])
            
            # 分离左右车道线
            left_lines = []
            right_lines = []
            
            for x1, y1, x2, y2 in adjusted_lines:
                # 计算线条斜率
                if x2 - x1 != 0:
                    slope = (y2 - y1) / (x2 - x1)
                    # 计算线条中点x坐标
                    mid_x = (x1 + x2) / 2
                    
                    # 根据斜率和位置判断是左线还是右线
                    if slope < -0.5 and mid_x < width / 2:
                        left_lines.append([x1, y1, x2, y2])
                    elif slope > 0.5 and mid_x > width / 2:
                        right_lines.append([x1, y1, x2, y2])
            
            # 计算平均左线和右线
            if left_lines:
                result['left_line'] = self._average_line(left_lines)
            
            if right_lines:
                result['right_line'] = self._average_line(right_lines)
            
            # 计算车道中心和宽度
            if result['left_line'] and result['right_line']:
                left_x = (result['left_line'][0] + result['left_line'][2]) / 2
                right_x = (result['right_line'][0] + result['right_line'][2]) / 2
                
                result['lane_width'] = right_x - left_x
                result['lane_center'] = (left_x + right_x) / 2 - width / 2
                result['lane_center'] /= width / 2  # 归一化
        
        return result
    
    def _average_line(self, lines):
        """
        计算一组线条的平均线条
        
        Args:
            lines: 线条列表，每条线表示为 [x1, y1, x2, y2]
            
        Returns:
            list: 平均线条 [x1, y1, x2, y2]
        """
        if not lines:
            return None
        
        # 计算所有点的平均位置
        points = []
        for line in lines:
            points.append((line[0], line[1]))
            points.append((line[2], line[3]))
        
        # 按y坐标排序，分为上下两组
        points.sort(key=lambda p: p[1])
        mid = len(points) // 2
        
        # 计算上半部分和下半部分的平均点
        upper_points = points[:mid]
        lower_points = points[mid:]
        
        if upper_points and lower_points:
            upper_avg_x = sum(p[0] for p in upper_points) / len(upper_points)
            upper_avg_y = sum(p[1] for p in upper_points) / len(upper_points)
            
            lower_avg_x = sum(p[0] for p in lower_points) / len(lower_points)
            lower_avg_y = sum(p[1] for p in lower_points) / len(lower_points)
            
            return [int(upper_avg_x), int(upper_avg_y), 
                   int(lower_avg_x), int(lower_avg_y)]
        
        return None
    
    def draw_line_following_results(self, image, line_result):
        """
        在图像上绘制巡线检测结果
        
        Args:
            image: 输入图像（BGR格式的numpy数组）
            line_result: detect_line_for_following的返回结果
            
        Returns:
            numpy.ndarray: 绘制了巡线检测结果的图像
        """
        # 创建图像副本
        result = image.copy()
        
        # 绘制检测到的线条
        for line in line_result['lines']:
            x1, y1, x2, y2 = line
            cv2.line(result, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # 绘制状态信息
        status_text = f"Status: {line_result['status']}"
        offset_text = f"Offset: {line_result['offset']:.2f}"
        angle_text = f"Angle: {line_result['angle']:.1f}°"
        
        cv2.putText(result, status_text, (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.putText(result, offset_text, (10, 60), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.putText(result, angle_text, (10, 90), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        # 绘制中心线
        height, width = image.shape[:2]
        cv2.line(result, (width // 2, 0), (width // 2, height), (255, 255, 0), 1)
        
        return result